'''
Split Dataset into Train & Valid & Test Dataset
'''
from pathlib import Path
import pandas as pd
from tqdm import tqdm
import random
import jsonlines


def read_jsonl(file_path):
    data_list = []
    with jsonlines.open(file_path) as reader:
        for obj in reader:
            data_list.append(obj)
    return data_list

def write_jsonl(data_list, save_path):
    with jsonlines.open(save_path, mode='w') as writer:
        for data in data_list:
            writer.write(data)

def split_data(data_list, ratio_list):
    random.shuffle(data_list)
    acc = 0
    ratio_acc = []
    for ratio in ratio_list:
        acc += ratio
        ratio_acc.append(acc)
    datasets = []

    ptr = 0
    for ratio in ratio_acc:
        # print(f"{ptr}: {int( len(data_list) * ratio ) }")
        datasets.append(data_list[ptr: int( len(data_list) * ratio )])
        ptr = int( len(data_list) * ratio )
    return datasets

def data_integrity(data_list):
    new_data_list = []
    broken_img_path = '/remote-home/share/medical/public/MedICaT/release/figures/ffd83b6453f94f2a1ddb346e324f5bdbf228c1f3_4-Figure3-1.png'
    root_dir = '/remote-home/share/medical/public/MedICaT'
    for obj in tqdm(data_list):
        img_path = f"{root_dir}/release/figures/{obj['pdf_hash']}_{obj['fig_uri']}"
        if img_path != broken_img_path:
            new_data_list.append(obj)
    return new_data_list

if __name__ == '__main__':
    print("\033[42mSplit Dataset\33[0m")
    dataset_dir = Path('/remote-home/share/medical/public/MedICaT')
    # data_path = dataset_dir / 'subfig_subcap_val.jsonl'
    data_path = dataset_dir / 'release/s2_full_figures_oa_nonroco_combined_medical_top4_public.jsonl'
    data_list = read_jsonl(file_path=data_path)
    print(f"total: {len(data_list)}")
    data_list = data_integrity(data_list)
    print(f"total: {len(data_list)}")

    trainset, validset, testset = split_data(data_list, [0.8, 0.1, 0.1])
    print( len(trainset), len(validset), len(testset) )
    write_jsonl(trainset, save_path='/remote-home/weixionglin/vlp/Analysis/medicat/train.jsonl')
    write_jsonl(validset, save_path='/remote-home/weixionglin/vlp/Analysis/medicat/valid.jsonl')
    write_jsonl(testset, save_path='/remote-home/weixionglin/vlp/Analysis/medicat/test.jsonl')
    testset_2 = read_jsonl('/remote-home/weixionglin/vlp/Analysis/medicat/test.jsonl')
    print(f"testset_2: {len(testset_2)}")
